{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "df630da9",
   "metadata": {},
   "source": [
    "# FNormalize\n",
    "\n",
    "FNormalize 是 TVM Relax 中的特殊属性，用于标准化算子的参数格式。\n",
    "\n",
    "测试主要验证 FNormalize 在不同场景下的行为，包括：\n",
    "1. 解析 TVMScript 时是否抑制 FNormalize 应用\n",
    "2. 在 C++ 和 Python 的 Mutator 中是否正确应用 FNormalize\n",
    "3. 标准化后的调用节点是否格式良好\n",
    "4. 未标准化的调用节点是否格式不良\n",
    "5. 针对特定算子（如 `call_tir`、`call_tir_inplace`、`call_tir_with_grad`）的参数元组内联处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c26f07fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "import tvm.testing\n",
    "import tvm.relax.testing.transform\n",
    "\n",
    "from tvm import relax\n",
    "from tvm.script import ir as I\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2a1bfdef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def register_custom_op(define_normalization=True):\n",
    "    \"\"\"\n",
    "    用于测试目的的自定义算子\n",
    "\n",
    "    这个自定义算子会忽略其第二个参数。如果没有可以忽略的第二个参数，\n",
    "    FNormalize 会添加额外的参数以便正确地忽略它。\n",
    "\n",
    "    参数:\n",
    "        define_normalization: 布尔值，决定是否为算子定义 FNormalize 属性\n",
    "\n",
    "    返回:\n",
    "        自定义算子实例\n",
    "    \"\"\"\n",
    "\n",
    "    op_name = \"custom_op.ignore_second_argument\"\n",
    "\n",
    "    def infer_struct_info(call: relax.Call, context: relax.BlockBuilder):\n",
    "        \"\"\"推断调用的结构信息，直接返回第一个参数的结构信息\"\"\"\n",
    "        return call.args[0].struct_info\n",
    "\n",
    "    def normalize(context: relax.BlockBuilder, call: relax.Call):\n",
    "        \"\"\"标准化调用参数：如果只有一个参数，则添加空元组作为第二个参数\"\"\"\n",
    "        if len(call.args) == 1:\n",
    "            return relax.Call(call.op, [call.args[0], relax.Tuple([])])\n",
    "        else:\n",
    "            return call\n",
    "\n",
    "    def legalize(context: relax.BlockBuilder, call: relax.Call):\n",
    "        \"\"\"合法化调用：直接返回第一个参数\"\"\"\n",
    "        return call.args[0]\n",
    "\n",
    "    # 定义算子属性\n",
    "    op_attrs = {\n",
    "        \"FInferStructInfo\": infer_struct_info,  # 推断结构信息\n",
    "        \"FLegalize\": legalize,  # 合法化处理\n",
    "        \"FPurity\": True,  # 标记为纯函数\n",
    "    }\n",
    "    if define_normalization:\n",
    "        op_attrs[\"FNormalize\"] = normalize  # 如果需要，添加标准化函数\n",
    "\n",
    "    # 注册算子属性\n",
    "    for key, value in op_attrs.items():\n",
    "        tvm.ir.register_op_attr(op_name, key, value)\n",
    "\n",
    "    op = tvm.ir.Op.get(op_name)\n",
    "    return op\n",
    "\n",
    "custom_op = register_custom_op(define_normalization=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98a16665",
   "metadata": {},
   "source": [
    "## 测试在解析 TVMScript 时是否抑制应用 FNormalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "40deea06",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">func</span>(A: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor:\n",
       "    gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor <span style=\"color: #A2F; font-weight: bold\">=</span> custom_op<span style=\"color: #A2F; font-weight: bold\">.</span>ignore_second_argument(A)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@R.function(check_well_formed=False)  # 禁用格式良好检查\n",
    "def func(A: R.Tensor) -> R.Tensor:\n",
    "    return relax.Call(custom_op, [A])\n",
    "func.show()\n",
    "# 提取函数体内的调用表达式\n",
    "call_expr = func.body.blocks[0].bindings[0].value\n",
    "assert isinstance(\n",
    "    call_expr, relax.Call\n",
    "), \"测试实现错误，未提取到正确的表达式\"\n",
    "assert (\n",
    "    len(call_expr.args) == 1\n",
    "), \"期望 TVMScript 抑制使用 FNormalize，按原样生成参数\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0599367",
   "metadata": {},
   "source": [
    "## 测试在 C++ 的 ExprMutator 子类中是否应用 FNormalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "28c2abe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module(check_well_formed=False)\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor):\n",
    "        return relax.Call(custom_op, [A])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "29da5e88",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module\n",
    "class Expected:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor):\n",
    "        return relax.Call(custom_op, [A, R.tuple()])\n",
    "\n",
    "# 应用空的 C++ Mutator\n",
    "After = tvm.relax.testing.transform.ApplyEmptyCppMutator()(Before)\n",
    "\n",
    "# 验证 Before 和 After 不相等，但 After 与 Expected 相等\n",
    "assert not tvm.ir.structural_equal(Before, After)\n",
    "tvm.ir.assert_structural_equal(Expected, After)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fedfc5e7",
   "metadata": {},
   "source": [
    "## 测试在 Python 的 ExprMutator 子类中是否应用 FNormalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2f2a60ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "@R.function(private=True, check_well_formed=False)\n",
    "def before(A: R.Tensor):\n",
    "    return relax.Call(custom_op, [A])\n",
    "\n",
    "@R.function(private=True)\n",
    "def expected(A: R.Tensor):\n",
    "    return relax.Call(custom_op, [A, R.tuple()])\n",
    "\n",
    "@relax.expr_functor.mutator\n",
    "class EmptyPyExprMutator(relax.PyExprMutator):\n",
    "    \"\"\"默认的 ExprMutator 实现\"\"\"\n",
    "    ...\n",
    "\n",
    "# 应用 Python Mutator\n",
    "after = EmptyPyExprMutator().visit_expr(before)\n",
    "\n",
    "# 验证 before 和 after 不相等，但 after 与 expected 相等\n",
    "assert not tvm.ir.structural_equal(before, after)\n",
    "tvm.ir.assert_structural_equal(expected, after)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92b54409",
   "metadata": {},
   "source": [
    "## 测试如果 FNormalize 不应用更改，IR 是否格式良好"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "97ac3e15",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module\n",
    "class Module:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor):\n",
    "        return relax.Call(custom_op, [A, A])\n",
    "\n",
    "# 验证模块格式良好\n",
    "assert relax.analysis.well_formed(Module)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3028ddc6",
   "metadata": {},
   "source": [
    "## 测试如果 FNormalize 应用更改，IR 是否格式不良"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4c552189",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[18:02:22] /media/pc/data/lxw/ai/tvm/src/relax/analysis/well_formed.cc:134: Warning: This IR is not well formed: If an operator defines an operator-specific normalization function (FNormalize), calls to that operator must be normalized with it.  However, normalization of custom_op.ignore_second_argument(A) resulted in custom_op.ignore_second_argument(A, R.tuple())\n"
     ]
    }
   ],
   "source": [
    "@I.ir_module(check_well_formed=False)\n",
    "class Module:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor):\n",
    "        return relax.Call(custom_op, [A])\n",
    "\n",
    "for define_normalization in [1, 0]:\n",
    "    for key in [\"FInferStructInfo\", \"FLegalize\", \"FPurity\", \"FNormalize\"]:\n",
    "        custom_op.reset_attr(key)\n",
    "    custom_op = register_custom_op(define_normalization=define_normalization)\n",
    "    # 如果定义了 FNormalize，则模块格式不良；否则格式良好\n",
    "    if define_normalization:\n",
    "        assert not relax.analysis.well_formed(Module)\n",
    "    else:\n",
    "        assert relax.analysis.well_formed(Module)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58c45859",
   "metadata": {},
   "source": [
    "## 测试 FNormalize 是否为 R.call_tir 内联参数元组"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1c11fc96",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module(check_well_formed=False)\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor([16], \"float32\")):\n",
    "        cls = Before\n",
    "        args = (A,)\n",
    "        return relax.Call(\n",
    "            tvm.ir.Op.get(\"relax.call_tir\"),\n",
    "            [cls.multiply_by_two, args],\n",
    "            sinfo_args=[A.struct_info],\n",
    "        )\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def multiply_by_two(A: T.Buffer(16, \"float32\"), B: T.Buffer(16, \"float32\")):\n",
    "        for i in range(16):\n",
    "            B[i] = A[i] * 2.0\n",
    "\n",
    "@I.ir_module\n",
    "class Expected:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor([16], \"float32\")):\n",
    "        cls = Expected\n",
    "        args = (A,)\n",
    "        return relax.Call(\n",
    "            tvm.ir.Op.get(\"relax.call_tir\"),\n",
    "            [cls.multiply_by_two, relax.Tuple([A])],  # 内联的元组\n",
    "            sinfo_args=[A.struct_info],\n",
    "        )\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def multiply_by_two(A: T.Buffer(16, \"float32\"), B: T.Buffer(16, \"float32\")):\n",
    "        for i in range(16):\n",
    "            B[i] = A[i] * 2.0\n",
    "\n",
    "# 应用 C++ Mutator\n",
    "After = tvm.relax.testing.transform.ApplyEmptyCppMutator()(Before)\n",
    "\n",
    "# 验证结果\n",
    "assert not tvm.ir.structural_equal(Before, After)\n",
    "tvm.ir.assert_structural_equal(Expected, After)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94c14246",
   "metadata": {},
   "source": [
    "## 测试 FNormalize 是否为 R.call_tir 内联参数元组"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9b36e4d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module(check_well_formed=False)\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(args: R.Tuple([R.Tensor([16], \"float32\")])):\n",
    "        cls = Before\n",
    "        return relax.Call(\n",
    "            tvm.ir.Op.get(\"relax.call_tir\"),\n",
    "            [cls.multiply_by_two, args],\n",
    "            sinfo_args=[args[0].struct_info],\n",
    "        )\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def multiply_by_two(A: T.Buffer(16, \"float32\"), B: T.Buffer(16, \"float32\")):\n",
    "        for i in range(16):\n",
    "            B[i] = A[i] * 2.0\n",
    "\n",
    "@I.ir_module\n",
    "class Expected:\n",
    "    @R.function\n",
    "    def main(args: R.Tuple([R.Tensor([16], \"float32\")])):\n",
    "        cls = Expected\n",
    "        return relax.Call(\n",
    "            tvm.ir.Op.get(\"relax.call_tir\"),\n",
    "            [cls.multiply_by_two, relax.Tuple([args[0]])],  # 从函数参数构建的内联元组\n",
    "            sinfo_args=[args[0].struct_info],\n",
    "        )\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def multiply_by_two(A: T.Buffer(16, \"float32\"), B: T.Buffer(16, \"float32\")):\n",
    "        for i in range(16):\n",
    "            B[i] = A[i] * 2.0\n",
    "\n",
    "# 应用 C++ Mutator\n",
    "After = tvm.relax.testing.transform.ApplyEmptyCppMutator()(Before)\n",
    "\n",
    "# 验证结果\n",
    "assert not tvm.ir.structural_equal(Before, After)\n",
    "tvm.ir.assert_structural_equal(Expected, After)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab12add0",
   "metadata": {},
   "source": [
    "## 测试 FNormalize 是否为 R.call_tir_inplace 内联参数元组"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "eaf47fd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CallTIRInplaceAttrs 在 Python API 中构造困难，因此先声明期望的模块并重用其属性\n",
    "@I.ir_module\n",
    "class Expected:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor([16], \"float32\")):\n",
    "        cls = Expected\n",
    "        args = (A,)\n",
    "        return R.call_tir_inplace(\n",
    "            cls.multiply_by_two,\n",
    "            A,\n",
    "            inplace_indices=[0],\n",
    "            out_sinfo=[A.struct_info],\n",
    "        )\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def multiply_by_two(A: T.Buffer(16, \"float32\")):\n",
    "        for i in range(16):\n",
    "            A[i] = A[i] * 2.0\n",
    "\n",
    "# 提取 inplace_attrs\n",
    "inplace_attrs = Expected[\"main\"].body.blocks[0].bindings[1].value.attrs\n",
    "\n",
    "@I.ir_module(check_well_formed=False)\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor([16], \"float32\")):\n",
    "        cls = Before\n",
    "        args = (A,)\n",
    "        return relax.Call(\n",
    "            tvm.ir.Op.get(\"relax.call_tir_inplace\"),\n",
    "            [cls.multiply_by_two, args],\n",
    "            attrs=inplace_attrs,\n",
    "            sinfo_args=[A.struct_info],\n",
    "        )\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def multiply_by_two(A: T.Buffer(16, \"float32\")):\n",
    "        for i in range(16):\n",
    "            A[i] = A[i] * 2.0\n",
    "\n",
    "# 应用 C++ Mutator\n",
    "After = tvm.relax.testing.transform.ApplyEmptyCppMutator()(Before)\n",
    "\n",
    "# 验证结果\n",
    "assert not tvm.ir.structural_equal(Before, After)\n",
    "tvm.ir.assert_structural_equal(Expected, After)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "703c0d63",
   "metadata": {},
   "source": [
    "## 测试 FNormalize 是否为 R.call_tir_with_grad 内联参数元组"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1e8134cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CallTIRWithGradAttrs 在 Python API 中构造困难，因此先声明期望的模块并重用其属性\n",
    "@I.ir_module\n",
    "class Expected:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor([16], \"float32\")):\n",
    "        cls = Expected\n",
    "        args = (A,)\n",
    "        return R.call_tir_with_grad(\n",
    "            cls.multiply_by_two,\n",
    "            A,\n",
    "            out_sinfo=[A.struct_info],\n",
    "            te_grad_name=\"f_grad\",\n",
    "        )\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def multiply_by_two(A: T.Buffer(16, \"float32\"), B: T.Buffer(16, \"float32\")):\n",
    "        for i in range(16):\n",
    "            B[i] = A[i] * 2.0\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def f_grad(\n",
    "        A: T.Buffer(16, \"float32\"), B: T.Buffer(16, \"float32\"), Grad: T.Buffer(16, \"float32\")\n",
    "    ):\n",
    "        for i in range(16):\n",
    "            Grad[i] = 2.0\n",
    "\n",
    "# 提取 with_grad_attrs\n",
    "with_grad_attrs = Expected[\"main\"].body.blocks[0].bindings[1].value.attrs\n",
    "\n",
    "@I.ir_module(check_well_formed=False)\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor([16], \"float32\")):\n",
    "        cls = Before\n",
    "        args = (A,)\n",
    "        return relax.Call(\n",
    "            tvm.ir.Op.get(\"relax.call_tir_with_grad\"),\n",
    "            [cls.multiply_by_two, args],\n",
    "            attrs=with_grad_attrs,\n",
    "            sinfo_args=[A.struct_info],\n",
    "        )\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def multiply_by_two(A: T.Buffer(16, \"float32\"), B: T.Buffer(16, \"float32\")):\n",
    "        for i in range(16):\n",
    "            B[i] = A[i] * 2.0\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def f_grad(\n",
    "        A: T.Buffer(16, \"float32\"), B: T.Buffer(16, \"float32\"), Grad: T.Buffer(16, \"float32\")\n",
    "    ):\n",
    "        for i in range(16):\n",
    "            Grad[i] = 2.0\n",
    "\n",
    "# 应用 C++ Mutator\n",
    "After = tvm.relax.testing.transform.ApplyEmptyCppMutator()(Before)\n",
    "\n",
    "# 验证结果\n",
    "assert not tvm.ir.structural_equal(Before, After)\n",
    "tvm.ir.assert_structural_equal(Expected, After)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "294e6236",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
